"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/4 20:23
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    generate_env_compare_pic.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    generate_env_compare_pic_VER0.1

------------------------------------------------------------------
# @Description:
    这里将从每一个测试环境项目中抽取一个ID的深度伪彩作为对比，
    实际拼接为 [color][orbbecDepth][rkDepthV0][rkDepthV1][....]

------------------------------------------------------------------
"""
import os
import cv2
import numpy as np
from tqdm import tqdm
from natsort import natsorted
from utils import Utils, Colorizer
from parser_cam_param import ParserCamParam


def get_all_env_test_disp_image(root_dir: str, key1='_disp', key2='.png'):
    if not os.path.exists(root_dir):
        print(' [ERROR] Test directory not exists!')
        return

    img_path_list = Utils().findFilesWithPathlibRecursive(root_dir, key1, key2)

    # skip precision category
    img_path_list = natsorted(img_path_list)
    img_path_list = [img_path for img_path in img_path_list if 'precision' not in img_path]

    return natsorted(img_path_list)


def get_rgb_and_orbbec_depth_img_path(env_disp_img_path: str, key='/rk_result_0708'):
    if not os.path.exists(env_disp_img_path):
        print(' [ERROR] Abnormal disp image path.')
        return

    rgb_img_path = env_disp_img_path.replace(key, '').replace('_disp', '_color')
    dp_img_path = rgb_img_path.replace('_color', '_depth')

    return rgb_img_path, dp_img_path


def generate_res_image(json_path: str, root_dir: str, disp_w, disp_h, key='/rk_result_0710'):
    if not os.path.exists(json_path):
        print(' [ERROR] Check json file path.')
        return

    if not os.path.exists(root_dir):
        print(' [ERROR] Test directory is empty')
        return

    # init
    colorize = Colorizer(levels=4096, depth_range=(0, 3000))
    parser = ParserCamParam(json_path)

    # get module param
    cam_param, ratio_w, ratio_h = parser.parser_cam_param(disp_w, disp_h)

    # get all disp image
    img_path_list = get_all_env_test_disp_image(root_dir)

    # gen depth image
    for idx, disp_path in enumerate(tqdm(img_path_list, position=0, desc=' Gen depth...', colour='cyan')):
        depth_path = disp_path.replace('_disp', '_depth')
        disp_img = cv2.imread(disp_path, cv2.IMREAD_UNCHANGED)
        dp_img = Utils().disp_to_depth(disp_img,
                                       cam_param['q_coeffs']['Q23'],
                                       cam_param['q_coeffs']['Q32'],
                                       cam_param['q_coeffs']['Q33'],
                                       sub_pixel_value=64,
                                       zoom_ratio=1.0,
                                       max_dis=3000)
        cv2.imwrite(depth_path, dp_img)
        pse_dp_img = colorize.colorize(dp_img, (0, 3000), bins=6000)
        pse_dp_img = cv2.resize(pse_dp_img, (640, 400))

        rgb_path, orbbec_dp_path = get_rgb_and_orbbec_depth_img_path(disp_path, key=key)
        rgb_img = cv2.resize(cv2.imread(rgb_path, cv2.IMREAD_COLOR), (711, 400))
        orbbec_dp_img = cv2.resize(cv2.imread(orbbec_dp_path, cv2.IMREAD_UNCHANGED), (640, 400))
        pse_orbbec_dp_img = colorize.colorize(orbbec_dp_img, (0, 3000), bins=6000)

        # 拼接 color pse_dp_img pse_orbbec_dp_img
        combined_img = np.concatenate((rgb_img, pse_orbbec_dp_img, pse_dp_img), axis=1)
        combined_img_path = disp_path.replace('_disp', '_combined').replace('.png', '.jpg')
        cv2.imwrite(combined_img_path, combined_img)


def generate_iter_compare_result(json_path: str, cur_ver_dir: str, disp_w, disp_h, last_ver_key='/rk_result_0708', cur_ver_key='/rk_result_0710'):
    if not os.path.exists(json_path):
        print(' [ERROR] Check json file path.')
        return

    if not os.path.exists(cur_ver_dir):
        print(' [ERROR] Test directory is empty')
        return

    # init
    colorize = Colorizer(levels=4096, depth_range=(0, 3000))
    parser = ParserCamParam(json_path)

    # get module param
    cam_param, ratio_w, ratio_h = parser.parser_cam_param(disp_w, disp_h)

    # get all disp image
    img_path_list = Utils().findFilesWithPathlibRecursive(cur_ver_dir, '_depth', '.png')
    img_path_list = [img_path for img_path in img_path_list if '_align' not in img_path]

    # skip precision category
    img_path_list = natsorted(img_path_list)

    for idx, depth_path in enumerate(tqdm(img_path_list, position=0, desc=' Colorizing...', colour='cyan')):
        # cur depth image.
        cur_dp_img = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
        cur_dp_img = np.clip(cur_dp_img, 10, 3000)
        cur_pse_img = colorize.colorize(cur_dp_img, (0, 3000), bins=6000)
        cur_pse_img = cv2.resize(cur_pse_img, (640, 400))

        # last depth image.
        last_dp_img_path = depth_path.replace(cur_ver_key, last_ver_key)
        last_dp_img = cv2.imread(last_dp_img_path, cv2.IMREAD_UNCHANGED)
        last_dp_img = np.clip(last_dp_img, 10, 3000)
        last_pse_img = colorize.colorize(last_dp_img, (0, 3000), bins=6000)
        last_pse_img = cv2.resize(last_pse_img, (640, 400))

        rgb_img_path = depth_path.replace(cur_ver_key, '').replace('_depth', '_color')
        rgb_img = cv2.resize(cv2.imread(rgb_img_path, cv2.IMREAD_COLOR), (711, 400))

        ir_img_path = depth_path.replace(cur_ver_key, '').replace('_depth', '_speckle-left')
        ir_img = cv2.imread(ir_img_path, cv2.IMREAD_UNCHANGED)
        ir_img = cv2.resize(ir_img, (640, 400))
        ir_img_rgb = np.stack((ir_img,) * 3, axis=-1)

        orbbec_dp_img_path = depth_path.replace(cur_ver_key, '')
        orbbec_dp_img = cv2.imread(orbbec_dp_img_path, cv2.IMREAD_UNCHANGED)
        orbbec_dp_img = np.clip(orbbec_dp_img, 10, 3000)
        orbbec_dp_img = cv2.resize(orbbec_dp_img, (640, 400))
        pse_orbbec_dp_img = colorize.colorize(orbbec_dp_img, (0, 3000), bins=6000)

        # 拼接 color pse_dp_img pse_orbbec_dp_img
        # combined_img = np.concatenate((rgb_img, pse_orbbec_dp_img, last_pse_img, cur_pse_img), axis=1)
        combined_img = np.concatenate((rgb_img, ir_img_rgb, pse_orbbec_dp_img, last_pse_img), axis=1)

        combined_img_path = depth_path.replace('_depth', '_combined').replace('.png', '.jpg')
        cv2.imwrite(combined_img_path, combined_img)


if __name__ == '__main__':
    test_env_dir = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/r1920x1080_d1280x800/rk_result_0719"
    json_file_path = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/rgb1920x1080_ir1280x800.json"
    disp_img_w = 1280
    disp_img_h = 800

    # generate_res_image(json_file_path, test_env_dir, disp_img_w, disp_img_h, key='/rk_result_0710')
    generate_iter_compare_result(json_file_path, test_env_dir, disp_img_w, disp_img_h, last_ver_key='/rk_result_0708', cur_ver_key='/rk_result_0710')
